import os
from dotenv import load_dotenv, find_dotenv
from langchain_community.chat_models import ChatZhipuAI 
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationChain

_ = load_dotenv(find_dotenv())
model = os.environ.get('model')
if model is None:
    raise ValueError("model is not set in the .env file")

llm = ChatZhipuAI(model=model,
                  temperature=0.9,              
    )

memory=ConversationBufferMemory()
conversation = ConversationChain(
    llm=llm, 
    #verbose=True, 
    memory=memory
)

ret = conversation.predict(input="我昨天给张三买了五本书。")
print(f"[Input]我昨天给张三买了五本书。\n[Output]{ret}")
ret = conversation.predict(input="我昨天给谁买书了？")
print(f"[Input]:我昨天给谁买书了？\n[Output]:{ret}")
ret = conversation.predict(input="我昨天买了几本书？")
print(f"[Input]:我昨天买了几本书？\n[Output]:{ret}")

#print(f"\nConversationBufferMemory:\n{memory}")